#include <iostream>
#include <stdlib.h>
#include <cmath>
#include <random>
#include <numeric>

#include "gtest/gtest.h"

#include "kpextVkContext.h"
#include "KpExtGP.h"
#include "kpextOpImgF32C2FFT.h"

#include "kpextCtxGlobal.h"

using namespace std;

const float AfterFFT_Real[] = {
    1920.f, -128.f, -128.f, -128.f, -128.f, -128.f, -128.f, -128.f,
    -128.f, -128.f, -128.f, -128.f, -128.f, -128.f, -128.f, -128.f,
    -643.49945499f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    0.f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    -309.01933598f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    0.f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    -191.56553762f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    0.f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    -128.f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    0.f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    -85.52686565f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    0.f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    -53.01933598f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    0.f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    -25.46078302f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    0.f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    0.f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    0.f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    25.46078302f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    0.f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    53.01933598f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    0.f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    85.52686565f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    0.f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    128.f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    0.f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    191.56553762f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    0.f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    309.01933598f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    0.f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    643.49945499f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f,
    0.f,    0.f,    0.f,    0.f, 0.f,    0.f,    0.f,    0.f
};

const float AfterFFT_Imag[] = {
1920.f,  643.49945499f,  309.01933598f,  191.56553762f,   128.f,   85.52686565f,   53.01933598f,   25.46078302f,
   0.f,  -25.46078302f,  -53.01933598f,  -85.52686565f,  -128.f, -191.56553762f, -309.01933598f, -643.49945499f,
-128.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
   0.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
-128.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
   0.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
-128.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
   0.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
-128.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
   0.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
-128.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
   0.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
-128.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
   0.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
-128.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
   0.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
-128.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
   0.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
-128.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
   0.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
-128.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
   0.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
-128.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
   0.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
-128.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
   0.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
-128.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
   0.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
-128.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
   0.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
-128.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f,
   0.f,    0.f,    0.f,    0.f,     0.f,    0.f,    0.f,    0.f
};

//  正向变换
TEST(OpImgF32C2FFTTest, FFT){
    ASSERT_NE(nullptr, GvkContextPtr);

    int cols = 16, rows = 16;

    auto cmdbuf = GvkContextPtr->cmdbuf();

    size_t size = static_cast<size_t>(cols * rows);
    auto stagebuf1 = GvkContextPtr->stagingBuf(size * 4 * 2);
    auto stagebuf2 = GvkContextPtr->stagingBuf(size * 4 * 2, false, true);
    auto inbuf = GvkContextPtr->deviceBuf(size * 4 * 2);
    auto tempbuf = GvkContextPtr->deviceBuf(size * 4 * 2);
    auto op = std::make_shared<kpext::OpImgF32C2FFT>(
                GvkContextPtr
                );

    ASSERT_NE(nullptr, cmdbuf);
    ASSERT_NE(nullptr, stagebuf1);
    ASSERT_NE(nullptr, stagebuf2);
    ASSERT_NE(nullptr, inbuf);
    ASSERT_NE(nullptr, tempbuf);
    ASSERT_NE(nullptr, op);


    typedef struct {float x; float y;} Complex;
//    std::vector<float> result_e;
//    result_e.resize(size);
    {
        auto data = stagebuf1->data<Complex>();
        for (size_t i = 0; i < size; ++i) {
            data[i].x = static_cast<float>(i % cols);
            data[i].y = static_cast<float>(i / cols);
        }
        stagebuf1->unmap();
    }


    kpext::CmdBaseBufferCpy bufcpy;
    kpext::CmdBaseBufferSet bufset;
    kpext::CmdBuffersMemBarrier barrierTx;
    kpext::CmdBuffersMemBarrier barrierRx;

    barrierTx.prepare(
                {inbuf, tempbuf},
                vk::PipelineStageFlagBits::eTransfer,
                vk::PipelineStageFlagBits::eComputeShader,
                vk::AccessFlagBits::eTransferWrite,
                vk::AccessFlagBits::eShaderRead
                );
    barrierRx.prepare(
                {inbuf, tempbuf},
                vk::PipelineStageFlagBits::eComputeShader,
                vk::PipelineStageFlagBits::eTransfer,
                vk::AccessFlagBits::eShaderWrite,
                vk::AccessFlagBits::eTransferRead
                );

    op->rebuild(cols, rows,
                inbuf, tempbuf,
                cmdbuf,
                0
                );
    auto outbuf = std::dynamic_pointer_cast<kpext::DeviceBuffer>(
                op->whichOutput(inbuf, tempbuf)
                );
    ASSERT_NE(nullptr, outbuf);

    cmdbuf->begin();
    bufcpy.prepare(inbuf, stagebuf1);
    cmdbuf->record(&bufcpy);
    cmdbuf->record(&barrierTx);

    cmdbuf->record(op.get());

    cmdbuf->end();
    cmdbuf->submit();
    cmdbuf->waitExec();


    if(outbuf == tempbuf){ // 交换
        outbuf = inbuf;
        inbuf = tempbuf;
        tempbuf = outbuf;
    }

    op->rebuild(cols, rows,
                inbuf, tempbuf,
                cmdbuf,
                1
                );

    outbuf = std::dynamic_pointer_cast<kpext::DeviceBuffer>(
                op->whichOutput(inbuf, tempbuf)
                );
    ASSERT_NE(nullptr, outbuf);

    cmdbuf->begin();
    cmdbuf->record(op.get());

    bufcpy.prepare(stagebuf2, outbuf);
    cmdbuf->record(&barrierRx);
    cmdbuf->record(&bufcpy);

    cmdbuf->end();
    cmdbuf->submit();
    cmdbuf->waitExec();


    {
        auto p2 = stagebuf2->data<float>();
        stagebuf2->invalidateMapped();
//            SPDLOG_INFO(__FUNCTION__ " output: [{}] ", stagebuf2->vector<float>());

        float maxerr = 0;
        for (int r = 0; r < rows; ++r) {
            for (int c = 0; c < cols; ++c) {
                int idx = c + r * cols;
                idx = idx * 2;
                float real = p2[idx];
                float imag = p2[idx + 1];
                idx = idx / 2;
                float err = std::abs(real - AfterFFT_Real[idx]);
                if(err > maxerr){
                    maxerr = err;
                }
                err = std::abs(imag - AfterFFT_Imag[idx]);
                if(err > maxerr){
                    maxerr = err;
                }
            }
//            auto sptr = p2 + r * cols * 2;
//            std::vector<float> rowdisp(sptr, sptr + cols * 2);
//            SPDLOG_INFO(" {} ", rowdisp);
        }
        stagebuf2->unmap();
        std::cout << "max error = " << maxerr << std::endl;
        EXPECT_LE(maxerr, 1e-4f);
    }
}

//  逆向变换
TEST(OpImgF32C2FFTTest, IFFT){
    ASSERT_NE(nullptr, GvkContextPtr);

    int cols = 16, rows = 16;

    auto cmdbuf = GvkContextPtr->cmdbuf();

    size_t size = static_cast<size_t>(cols * rows);
    auto stagebuf1 = GvkContextPtr->stagingBuf(size * 4 * 2);
    auto stagebuf2 = GvkContextPtr->stagingBuf(size * 4 * 2, false, true);
    auto inbuf = GvkContextPtr->deviceBuf(size * 4 * 2);
    auto tempbuf = GvkContextPtr->deviceBuf(size * 4 * 2);
    auto op = std::make_shared<kpext::OpImgF32C2FFT>(
                GvkContextPtr
                );

    ASSERT_NE(nullptr, cmdbuf);
    ASSERT_NE(nullptr, stagebuf1);
    ASSERT_NE(nullptr, stagebuf2);
    ASSERT_NE(nullptr, inbuf);
    ASSERT_NE(nullptr, tempbuf);
    ASSERT_NE(nullptr, op);


    typedef struct {float x; float y;} Complex;
    std::vector<Complex> result_e;
    result_e.resize(size);
    {
        auto data = stagebuf1->data<Complex>();
        for (size_t i = 0; i < size; ++i) {
            result_e[i].x = static_cast<float>(i % cols);
            result_e[i].y = static_cast<float>(i / cols);
            data[i].x = AfterFFT_Real[i];
            data[i].y = AfterFFT_Imag[i];
        }
        stagebuf1->unmap();
    }


    kpext::CmdBaseBufferCpy bufcpy;
    kpext::CmdBaseBufferSet bufset;
    kpext::CmdBuffersMemBarrier barrierTx;
    kpext::CmdBuffersMemBarrier barrierRx;

    barrierTx.prepare(
                {inbuf, tempbuf},
                vk::PipelineStageFlagBits::eTransfer,
                vk::PipelineStageFlagBits::eComputeShader,
                vk::AccessFlagBits::eTransferWrite,
                vk::AccessFlagBits::eShaderRead
                );
    barrierRx.prepare(
                {inbuf, tempbuf},
                vk::PipelineStageFlagBits::eComputeShader,
                vk::PipelineStageFlagBits::eTransfer,
                vk::AccessFlagBits::eShaderWrite,
                vk::AccessFlagBits::eTransferRead
                );

    op->rebuild(cols, rows,
                inbuf, tempbuf,
                cmdbuf,
                1,
                1.0f/static_cast<float>(cols),
                true
                );
    auto outbuf = std::dynamic_pointer_cast<kpext::DeviceBuffer>(
                op->whichOutput(inbuf, tempbuf)
                );
    ASSERT_NE(nullptr, outbuf);

    cmdbuf->begin();
    bufcpy.prepare(inbuf, stagebuf1);
    cmdbuf->record(&bufcpy);
    cmdbuf->record(&barrierTx);

    cmdbuf->record(op.get());

    cmdbuf->end();
    cmdbuf->submit();
    cmdbuf->waitExec();


    if(outbuf == tempbuf){ // 交换
        outbuf = inbuf;
        inbuf = tempbuf;
        tempbuf = outbuf;
    }

    op->rebuild(cols, rows,
                inbuf, tempbuf,
                cmdbuf,
                0,
                1.0f/static_cast<float>(cols),
                true
                );

    outbuf = std::dynamic_pointer_cast<kpext::DeviceBuffer>(
                op->whichOutput(inbuf, tempbuf)
                );
    ASSERT_NE(nullptr, outbuf);

    cmdbuf->begin();
    cmdbuf->record(op.get());

    bufcpy.prepare(stagebuf2, outbuf);
    cmdbuf->record(&barrierRx);
    cmdbuf->record(&bufcpy);

    cmdbuf->end();
    cmdbuf->submit();
    cmdbuf->waitExec();


    {
        auto p2 = stagebuf2->data<float>();
        stagebuf2->invalidateMapped();
//            SPDLOG_INFO(__FUNCTION__ " output: [{}] ", stagebuf2->vector<float>());

        float maxerr = 0;
        for (int r = 0; r < rows; ++r) {
            for (int c = 0; c < cols; ++c) {
                int idx = c + r * cols;
                idx = idx * 2;
                float real = p2[idx];
                float imag = p2[idx + 1];
                idx = idx / 2;
                float err = std::abs(real - result_e[idx].x);
                if(err > maxerr){
                    maxerr = err;
                }
                err = std::abs(imag - result_e[idx].y);
                if(err > maxerr){
                    maxerr = err;
                }
            }
//            auto sptr = p2 + r * cols * 2;
//            std::vector<float> rowdisp(sptr, sptr + cols * 2);
//            SPDLOG_INFO(" {} ", rowdisp);
        }
        stagebuf2->unmap();
        std::cout << "max error = " << maxerr << std::endl;
        EXPECT_LE(maxerr, 1e-4f);
    }
}

//  正向变换用时测试
TEST(OpImgF32C2FFTTest, FFT_Benchmark){
    ASSERT_NE(nullptr, GvkContextPtr);

    int cols = 2048, rows = 1000;
    int cycletime = 500;

    auto cmdbuf = GvkContextPtr->cmdbuf();

    size_t size = static_cast<size_t>(cols * rows);
    auto stagebuf1 = GvkContextPtr->stagingBuf(size * 4 * 2);
    auto stagebuf2 = GvkContextPtr->stagingBuf(size * 4 * 2, false, true);
    auto inbuf = GvkContextPtr->deviceBuf(size * 4 * 2);
    auto tempbuf = GvkContextPtr->deviceBuf(size * 4 * 2);
    auto op = std::make_shared<kpext::OpImgF32C2FFT>(
                GvkContextPtr
                );

    ASSERT_NE(nullptr, cmdbuf);
    ASSERT_NE(nullptr, stagebuf1);
    ASSERT_NE(nullptr, stagebuf2);
    ASSERT_NE(nullptr, inbuf);
    ASSERT_NE(nullptr, tempbuf);
    ASSERT_NE(nullptr, op);


    typedef struct {float x; float y;} Complex;

    {
        std::random_device r;

        std::default_random_engine e1(r());
        std::normal_distribution<float> normal_dist(
                    0,  // mean
                    30000  // sigma
                    );

        auto data = stagebuf1->data<Complex>();
        for (size_t i = 0; i < size; ++i) {
            data[i].x = normal_dist(e1);
            data[i].y = normal_dist(e1);
        }
        stagebuf1->unmap();
    }


    kpext::CmdBaseBufferCpy bufcpy;
    kpext::CmdBaseBufferSet bufset;
    kpext::CmdBuffersMemBarrier barrierTx;
    kpext::CmdBuffersMemBarrier barrierRx;

    barrierTx.prepare(
                {inbuf, tempbuf},
                vk::PipelineStageFlagBits::eTransfer,
                vk::PipelineStageFlagBits::eComputeShader,
                vk::AccessFlagBits::eTransferWrite,
                vk::AccessFlagBits::eShaderRead
                );
    barrierRx.prepare(
                {inbuf, tempbuf},
                vk::PipelineStageFlagBits::eComputeShader,
                vk::PipelineStageFlagBits::eTransfer,
                vk::AccessFlagBits::eShaderWrite,
                vk::AccessFlagBits::eTransferRead
                );

    op->rebuild(cols, rows,
                inbuf, tempbuf,
                cmdbuf,
                0
                );
    auto outbuf = std::dynamic_pointer_cast<kpext::DeviceBuffer>(
                op->whichOutput(inbuf, tempbuf)
                );
    ASSERT_NE(nullptr, outbuf);

    cmdbuf->begin();
    bufcpy.prepare(inbuf, stagebuf1);
    cmdbuf->record(&bufcpy);
    cmdbuf->record(&barrierTx);

    cmdbuf->record(op.get());

    bufcpy.prepare(stagebuf2, outbuf);
    cmdbuf->record(&barrierRx);
    cmdbuf->record(&bufcpy);

    cmdbuf->end();

    //  重复发送命令执行， 计时
    fmt::print("\n Doing FFT benchmark test, please wait!...................\n");
    kpextSetLoggerLevel(KpExtLOG_LEVEL_ERROR);
    auto begin = std::chrono::high_resolution_clock::now();

    for (int i = 0; i < cycletime; ++i) {
        cmdbuf->submit();
        cmdbuf->waitExec();
    }

    auto end = std::chrono::high_resolution_clock::now();
    auto elapsed = std::chrono::duration_cast<std::chrono::nanoseconds>(end - begin);
    double pt = static_cast<double>(elapsed.count());

    pt = pt * 1e-6;
    fmt::print("\n *********************************************\n");
    fmt::print("OpImgF32C2FFT fft by row direction --> \n"
               "Frame size : ({}x{}),\n"
               "{} cycle,\n"
               "total time : {}ms,\n"
               "each FFT time : {}ms",
               rows, cols,
               cycletime, pt,
               pt / cycletime
               );
    fmt::print("\n *********************************************\n");

    kpextSetLoggerLevel(KpExtLOG_LEVEL_INFO);
    {
//        auto p2 = stagebuf2->data<float>();
//        stagebuf2->invalidateMapped();
//            SPDLOG_INFO(__FUNCTION__ " output: [{}] ", stagebuf2->vector<float>());

//        for (int r = 0; r < rows; ++r) {
//            for (int c = 0; c < cols; ++c) {
//                int idx = c + r * cols;
//            }
//            auto sptr = p2 + r * cols * 2;
//            std::vector<float> rowdisp(sptr, sptr + cols * 2);
//            SPDLOG_INFO(" {} ", rowdisp);
//        }
//        stagebuf2->unmap();
    }
}

GTEST_API_ int main(int argc, char **argv) {
    printf("Running main() from %s\n", __FILE__);
    //    kpextSetLoggerLevel(KpExtLOG_LEVEL_DEBUG);
    kpextSetLoggerLevel(KpExtLOG_LEVEL_INFO);
    //  kpextSetLoggerPattern("[%D %H:%M:%S.%e] [%^%l%$] %v");
    kpextSetLoggerPattern("   [SPDLOG_%l] %v");

    testing::InitGoogleTest(&argc, argv);

    testing::Environment* const foo_env =
            testing::AddGlobalTestEnvironment(new KpExtCtxEnv);
    UNUSE_VAR(foo_env);

    return RUN_ALL_TESTS();
}
